from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import TimeSeriesSplit
import pandas as pd
import numpy as np
import os, joblib

# 数据处理器
class DataProcessor:
    def __init__(self, args):
        self.args = args
        self.scaler = MinMaxScaler(feature_range=(-1, 1))
        
    def load_data(self):
        df = pd.read_csv(self.args.data_path, parse_dates=['Month'], index_col='Month')
        values = df['Passengers'].values.astype('float32')
        return values
    
    def split_data(self, values):
        split_index = int(len(values) * self.args.split_ratio)
        train = values[:split_index]
        test = values[split_index:]
        return train, test
    
    def create_sequences(self, data):
        X, y = [], []
        for i in range(len(data)-self.args.seq_length):
            X.append(data[i:i+self.args.seq_length])
            y.append(data[i+self.args.seq_length])
        return np.array(X), np.array(y)
    
    def scale_data(self, train, test=None):
        train_scaled = self.scaler.fit_transform(train.reshape(-1, 1))
        if test is not None:
            test_scaled = self.scaler.transform(test.reshape(-1, 1))
            return train_scaled.flatten(), test_scaled.flatten()
        return train_scaled.flatten()
    
    def save_scaler(self, save_dir):
            """保存scaler到指定目录"""
            os.makedirs(save_dir, exist_ok=True)
            scaler_path = os.path.join(save_dir, 'scaler.pkl')
            joblib.dump(self.scaler, scaler_path)
            return scaler_path
        
    @staticmethod
    def load_scaler(load_dir):
        """从指定目录加载scaler"""
        scaler_path = os.path.join(load_dir, 'scaler.pkl')
        return joblib.load(scaler_path)